from pathlib import Path
import numpy as np
from tqdm import tqdm
import argparse
import json

from utils import process_record, save_record, get_kfold, generate_stratified_split, get_split_distribution


DATA_ROOT = Path('/data/IDLab/DigiHealth/mit-bih-arrhythmia-database-1.0.0')
LABEL_MAPPING = {
    "N": 0, "L": 0, "R": 0, "e": 0, "j": 0,  # N
    "A": 1, "a": 1, "S": 1, "J": 1,  # SVEB
    "V": 2, "E": 2,  # VEB
    "F": 3,  # F
    "/": 4, "f": 4, "Q": 4  # Q
}
DS1_IDS = [
    101, 106, 108, 109, 112,
    114, 115, 116, 118, 119,
    122, 124, 201, 203, 205,
    207, 208, 209, 215, 220,
    223, 230
]
DS2_IDS = [
    100, 103, 105, 111, 113,
    117, 121, 123, 200, 202,
    210, 212, 213, 214, 219,
    221, 222, 228, 231, 232,
    233, 234
]

WINDOW_SIZE = 700 # in # of samples
SAMPLING_RATE = 360
NEW_SAMPLING_RATE = None
TIME_ENCODED = False
OUTPUT_FOLDER = Path(f'/data/IDLab/DigiHealth/processed_data/beat-to-beat/mit-bih-time/')
FOLDS_FOLDER = Path(f'/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/processed_data/configs')

TEST_RATIO = 1 / 8

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--stratified', action='store_true')
    parser.add_argument('--standard', action='store_true')
    parser.add_argument('--reversed', action='store_true')
    parser.add_argument('--kfold', action='store_true')
    args = parser.parse_args()

    if args.stratified:
        training_ids, testing_ids = generate_stratified_split(OUTPUT_FOLDER, testing_fraction=TEST_RATIO)
        training_distribution, testing_distribution = get_split_distribution(OUTPUT_FOLDER, training_ids.tolist(), testing_ids.tolist())
        json_dict = [{
            'training': {
                'distribution': training_distribution,
                'patient_ids': training_ids.tolist()
            },
            'testing': {
                'distribution': testing_distribution,
                'patient_ids': testing_ids.tolist()
            }
        },]
        print(json.dumps(json_dict, sort_keys=True, indent=4), file=open(FOLDS_FOLDER / f'mit_stratified_{int(1 / TEST_RATIO)}.json', 'w'))
    elif args.standard:
        training_ids = np.array(DS1_IDS)
        testing_ids = np.array(DS2_IDS)
        training_distribution, testing_distribution = get_split_distribution(OUTPUT_FOLDER, training_ids.tolist(), testing_ids.tolist())
        json_dict = [{
            'training': {
                'distribution': training_distribution,
                'patient_ids': training_ids.tolist()
            },
            'testing': {
                'distribution': testing_distribution,
                'patient_ids': testing_ids.tolist()
            }
        },]
        print(json.dumps(json_dict, sort_keys=True, indent=4), file=open(FOLDS_FOLDER / f'mit_standard.json', 'w'))
    elif args.reversed:
        training_ids = np.array(DS2_IDS)
        testing_ids = np.array(DS1_IDS)
        training_distribution, testing_distribution = get_split_distribution(OUTPUT_FOLDER, training_ids.tolist(), testing_ids.tolist())
        json_dict = [{
            'training': {
                'distribution': training_distribution,
                'patient_ids': training_ids.tolist()
            },
            'testing': {
                'distribution': testing_distribution,
                'patient_ids': testing_ids.tolist()
            }
        },]
        print(json.dumps(json_dict, sort_keys=True, indent=4), file=open(FOLDS_FOLDER / f'mit_reversed.json', 'w'))
    elif args.kfold:
        patient_ids = set([int(f.stem.split('_')[-1]) for f in OUTPUT_FOLDER.glob('PATIENT*')])
        patient_ids = np.array(list(patient_ids))
        folds = get_kfold(OUTPUT_FOLDER, patient_ids)
        print(json.dumps(folds, sort_keys=True, indent=4), file=open(FOLDS_FOLDER / 'mit_5_fold.json', 'w'))
    else:
        # for sr in range(100, 501, 25):
        if True:
            # OUTPUT_FOLDER = Path(f'/data/IDLab/DigiHealth/processed_data/beat-to-beat/mit-bih-time-{sr}/')
            print(OUTPUT_FOLDER.name)
            
            patient_ids = [int(p.stem) for p in DATA_ROOT.glob('*.hea')]
            for patient_id in tqdm(patient_ids):
                result = process_record(DATA_ROOT, patient_id, LABEL_MAPPING,
                                        window_size=WINDOW_SIZE,
                                        sampling_rate=SAMPLING_RATE, resampling=None,
                                        time_encoded=TIME_ENCODED)
                if result is not None:
                    patient_folder = OUTPUT_FOLDER / f"PATIENT_{patient_id}" / 'RECORD_0'
                    save_record(patient_folder, *result)
            

    
    